// app/controller/users.js
const Controller = require('egg').Controller;

const { PrismaClient } = require('@prisma/client');
const { PDFLoader } = require('langchain/document_loaders/fs/pdf');
const { TextLoader } = require('langchain/document_loaders/fs/text');
const { JSONLoader } = require('langchain/document_loaders/fs/json');
const { DocxLoader } = require('langchain/document_loaders/fs/docx');
const { RecursiveCharacterTextSplitter } = require('langchain/text_splitter');
const { OllamaEmbeddings } = require('@langchain/community/embeddings/ollama');
const { Chroma } = require('@langchain/community/vectorstores/chroma');
const fs = require('fs');
const path = require('path');
// // 注意删除临时文件
// await fs.unlink(file.filepath);

function toInt(str) {
  if (typeof str === 'number') return str;
  if (!str) return str;
  return parseInt(str, 10) || 0;
}

// 定义函数来处理转换
function convertToNewFormat(nodeFileData, target) {
  // 读取文件数据为 Buffer
  const data = fs.readFileSync(target);

  // 构建新对象
  const newData = {
    name: nodeFileData?.fieldname, // 按需求命名
    filename: nodeFileData.filename,
    type: nodeFileData.mime || nodeFileData.mimeType, // 可以根据需求选择其中一个字段作为文件类型
    data, // 文件数据的 Buffer
  };

  return newData;
}

class UserController extends Controller {
  constructor(props) {
    super(props);
    this.prisma = new PrismaClient();
  }

  async loadDocuments(file) {
    const Loaders = {
      pdf: PDFLoader,
      json: JSONLoader,
      docx: DocxLoader,
      doc: DocxLoader,
      txt: TextLoader,
      md: TextLoader,
    };
    const ext = (file.filename?.match(/\.(\w+)$/)?.[1] || 'txt').toLowerCase();
    if (!Loaders[ext]) {
      return {
        code: -1,
        data: `Unsupported file type: ${ext}`,
      };
    }
    const blob = new Blob([ file.data ], { type: file.type });
    return new Loaders[ext](blob).load();
  }

  async ingestDocument(file, collectionName, embedding, ollamaHost) {
    const docs = await this.loadDocuments(file);
    if (docs?.code === -1) {
      return { file, collectionName, embedding, ollamaHost };
    }

    const textSplitter = new RecursiveCharacterTextSplitter({
      chunkSize: 1000,
      chunkOverlap: 200,
    });
    const splits = await textSplitter.splitDocuments(docs);
    const embeddings = new OllamaEmbeddings({
      model: embedding,
      baseUrl: ollamaHost,
    });

    // Chromadb数据库设置--------------------------------
    const dbConfig = {
      collectionName,
      url: process.env.CHROMADB_URL,
    };
    const existingCollection = await Chroma.fromExistingCollection(
      embeddings,
      dbConfig
    );
    if (existingCollection) {
      await existingCollection.addDocuments(splits);
      console.log(`Chroma collection ${collectionName} updated`);
    } else {
      await Chroma.fromDocuments(splits, embeddings, dbConfig);
      console.log(`Chroma collection ${collectionName} created`);
    }
  }

  // get
  async index() {
    const ctx = this.ctx;
    const res = await this.prisma.knowledgeBase.findMany({
      include: {
        files: true,
      },
    });
    ctx.body = {
      code: 0,
      data: res,
      success: true,
    };
  }

  // get /:id
  async show() {
    const ctx = this.ctx;
    // ctx.params.id
    const knowledgeBase = await this.prisma.knowledgeBase.findUnique({
      where: {
        id: toInt(ctx.params.id),
      },
    });

    ctx.body = {
      code: 0,
      data: knowledgeBase,
      success: true,
    };
  }

  // post
  async create() {
    const ctx = this.ctx;
    const { name, description, embedding } = ctx.request.body;
    // 获取配置服务地址
    const { HOST } = process.env;
    const uploadedFiles = ctx.request.files;

    const affected = await this.prisma.knowledgeBase.create({
      data: {
        name,
        description,
        embedding,
        created: new Date(),
      },
    });
    console.log(`Created knowledge base ${name}: ${JSON.stringify(affected)}`);

    if (uploadedFiles.length > 0) {
      for (const uploadedFile of uploadedFiles) {
        // 同步读取文件内容 文件copy到项目下 方便查询
        const fileinfo = fs.readFileSync(uploadedFile.filepath);
        const arr = uploadedFile.filename.split('.');
        const filename =
          arr[0] + new Date().getTime() + '.' + arr[arr.length - 1];
        const target = path.join('app/public/uploads', filename);

        try {
          fs.writeFileSync(target, fileinfo);
        } catch (error) {
          throw error;
        } finally {
          await fs.promises.unlink(uploadedFile.filepath, err => {
            if (err) {
              throw err;
            }
          });
        }
        const url = `/public/uploads/${filename}`;

        // 转换buffer文件
        const file = convertToNewFormat(uploadedFile, target);

        await this.ingestDocument(
          file,
          `collection_${affected.id}`,
          affected.embedding,
          HOST
        );

        const createdKnowledgeBaseFile =
          await this.prisma.knowledgeBaseFile.create({
            data: {
              url,
              name: uploadedFile.filename,
              knowledgeBaseId: affected.id,
            },
          });

        console.log(createdKnowledgeBaseFile);
      }
    }

    ctx.status = 200;
    ctx.body = {
      code: 0,
      status: 'success',
    };
  }

  async update() {
    const ctx = this.ctx;
    const id = toInt(ctx.params.id);
    ctx.body = {
      status: 'success',
      id,
    };
  }

  async destroy() {
    const ctx = this.ctx;
    const id = toInt(ctx.params.id);
    const deletedKnowledgeBase = await this.prisma.knowledgeBase.delete({
      where: {
        id,
      },
    });

    ctx.status = 200;
    ctx.body = {
      code: 0,
      data: deletedKnowledgeBase,
      status: 'success',
    };
  }
}

module.exports = UserController;
